{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c2107c85-1c88-43be-ae3a-3a03cddca5d6",
   "metadata": {},
   "source": [
    "# Pymol Visualization\n",
    "## Objective\n",
    "\n",
    "The objective of this tutorial is to learn how to fully utilize `pymol_visualize()` method.\n",
    "\n",
    "This workflow allows you to:\n",
    "\n",
    "- Generate Python script for PyMOL visualization\n",
    "- Visualize different regions using PyMOL scene function\n",
    "- Create separate directories when making multiple visualization script\n",
    "\n",
    "In this specific example, we are using ketosteroid isomerase (KSI) as the model system. The structure for KSI is obtained from the PDB [1OH0](https://doi.org/10.2210/pdb1OH0/pdb) and MM-minimized prior to this tutorial.\n",
    "\n",
    "## Classes used in this example\n",
    "\n",
    "- [GenerateModel](https://qmzyme.readthedocs.io/en/latest/API/QMzyme.GenerateModel.html)\n",
    "- [QM_Method](https://qmzyme.readthedocs.io/en/latest/API/QMzyme.CalculateModel.html#qm-treatment)\n",
    "- [CA_terminal TruncationScheme](https://qmzyme.readthedocs.io/en/latest/API/QMzyme.TruncationSchemes.html#QMzyme.TruncationSchemes.CA_terminal)\n",
    "- [DistanceCutoff SelectionScheme](https://qmzyme.readthedocs.io/en/latest/API/QMzyme.SelectionSchemes.html#QMzyme.SelectionSchemes.DistanceCutoff)\n",
    "\n",
    "## Required Files\n",
    "\n",
    "To start, you will need:\n",
    "\n",
    "- A fully prepped and protonated PDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0717964a-60b3-44a8-a58b-fdff48afcfa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here are the necesary imports for this tutorial!\n",
    "\n",
    "import QMzyme \n",
    "from QMzyme.data import PDB\n",
    "from QMzyme.SelectionSchemes import DistanceCutoff"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "330922e7-e709-4238-9ce0-519a2e2f30f3",
   "metadata": {},
   "source": [
    "## Visualizing the Full Model\n",
    "\n",
    "To start with, let's initialize QMzymeModel and simply visualize the loaded structure file. This can be done by initializing QMzymeModel using `GenerateModel()` and using `pymol_visualize()` on the model to get python script for visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "37ca30ff-1dce-430c-a4db-54b6c806576f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Charge information not present. QMzyme will try to guess region charges based on residue names consistent with AMBER naming conventions (i.e., aspartate: ASP --> Charge: -1, aspartic acid: ASH --> Charge: 0.). See QMzyme.data.residue_charges for the full set.\n"
     ]
    }
   ],
   "source": [
    "model = QMzyme.GenerateModel(PDB)\n",
    "model.pymol_visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d76a7d4a-1249-4ada-93bb-00ac1c5a20e4",
   "metadata": {},
   "source": [
    "Now, when using this, you might have realized that the structure itself is hidden in the front interface. However, when you navigate towards `SCENES` tab in PyMOL, you can see that we see our starting structure!\n",
    "\n",
    "To organize things better, you can specify the output_folder. Using this will create a directory with the specified name, and will store the Python script and structure files needed to make the PyMOL visualization!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2afcfad1-8d5a-4634-94b0-690ae1aee412",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is done simply by adding an argument, output_folder=\"str\"\n",
    "model.pymol_visualize(output_dir=\"full model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03d6f1f0-ab2a-4d40-aed9-4d49ece84808",
   "metadata": {},
   "source": [
    "## Visualizing the QMzyme-prepared structure\n",
    "\n",
    "In addition to viewing the full model, `pymol_visualize()` can be used to examine the construction history of the region. Let's create a simple workflow using QMzyme, and use `pymol_visualize()` to see the QM region used for input file generation!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc072238-eb45-44ff-845c-b244c1ffdf2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We first initilize QMzymeModel by loading structure file.\n",
    "# The initial pdb file should be prepared prior (hydrogens must be present).\n",
    "model = QMzyme.GenerateModel(PDB)\n",
    "\n",
    "# This adds unknown residue charge.\n",
    "QMzyme.data.residue_charges.update({'EQU': -1})\n",
    "\n",
    "# This is used to set catalytic center.\n",
    "model.set_catalytic_center(selection='resid 263')\n",
    "\n",
    "# Use DistanceCutoff to create 3 Å region around catalytic center.\n",
    "model.set_region(selection=DistanceCutoff, cutoff=3)\n",
    "\n",
    "# Selecting and freeznig CA atoms.\n",
    "c_alpha_atoms = model.cutoff_3.get_atoms(attribute='name', value='CA')\n",
    "model.cutoff_3.set_fixed_atoms(atoms=c_alpha_atoms)\n",
    "\n",
    "# Set method to a region.\n",
    "qm_method = QMzyme.QM_Method(\n",
    "    basis_set='6-31G*', \n",
    "    functional='wB97X-D3', \n",
    "    qm_input='OPT FREQ', \n",
    "    program='orca'\n",
    ")\n",
    "qm_method.assign_to_region(region=model.cutoff_3)\n",
    "\n",
    "# Truncate the model\n",
    "model.truncate()\n",
    "\n",
    "# Write QM input file\n",
    "model.write_input()\n",
    "\n",
    "# Create Python script for visualization\n",
    "model.pymol_visualize(output_dir=\"cutoff_3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cd5e409-83bd-4341-b415-a8cdb1cb18fe",
   "metadata": {},
   "source": [
    "## Visualizing the .pkl file\n",
    "\n",
    "After running a full workflow, QMzyme will create a pickle file, which will contain the construction history of the QM input file. The pickle file can be loaded to recreate the QMzymeModel for subsequent structural analysis!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9ec28cb1-3efc-4e1a-a0e7-bafe23cf0f53",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = QMzyme.GenerateModel(pickle_file=f'1oh0.pkl')\n",
    "model.pymol_visualize(output_dir=\"pickled_model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "611ad6d9-79b6-471f-83e3-2a723e52ed77",
   "metadata": {},
   "source": [
    "## References Cited\n",
    "\n",
    "The PyMOL Molecular Graphics System, Version 3.0 Schrödinger, LLC."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}